package cn.doitedu.ml.demo

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.SparkSession

import scala.collection.mutable

object NaiveBayesDemoTrainner {
  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder().appName("朴素贝叶斯分类算法示例：出轨预测").master("local").getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._


    // 加载样本数据集
    val sample = spark.read.option("header","true").csv("user_portrait/data/chugui/sample")

    /**
     * 特征工程
     */
    // 特征值数字化
    val sampleDatumn = sample.selectExpr(
      "name",
      "case when job='老师' then 1.0 when job='公务员' then 2.0   else 3.0 end as job",
      "case when income='低' then 1.0 when income='中' then 2.0   else 3.0 end as income",
      "case when age='中年' then 1.0 when age='青年' then 2.0   else 3.0 end as age",
      "case when sex='男' then 1.0  else 2.0 end as sex",
      "case when label='出轨' then 1.0  else 0.0 end as label"
    )

    // 特征向量化
    val to_vec = udf((arr:mutable.WrappedArray[Double])=>{
      Vectors.dense(arr.toArray)
    })

    val sampleVec = sampleDatumn.select(to_vec(array('job,'income,'age,'sex)) as "vec",'label)

    /**
     * 训练朴素贝叶斯算法模型
     */
    val naiveBayes = new NaiveBayes()
        .setFeaturesCol("vec")
        .setLabelCol("label")
        .setSmoothing(0.01)
    val model: NaiveBayesModel = naiveBayes.fit(sampleVec)

    // 保存训练好的模型
    model.save("user_portrait/data/chugui/model")



    spark.close()
  }

}
